package top.cyuw.simplerpc.registry.memory;

import top.cyuw.simplerpc.URL;
import top.cyuw.simplerpc.dto.RpcRequest;
import top.cyuw.simplerpc.extension.ExtensionLoader;
import top.cyuw.simplerpc.loadbalance.LoadBalance;
import top.cyuw.simplerpc.loadbalance.LoadBalanceSelector;
import top.cyuw.simplerpc.registry.ServiceDiscovery;
import top.cyuw.simplerpc.util.LockUtil;

import java.net.InetSocketAddress;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author chen
 * @date 2023/3/16 21:10
 */
public class MemoryServiceDiscovery implements ServiceDiscovery {

    protected final Map<String, List<URL>> addressMap;
    private final LoadBalance loadBalance;

    public MemoryServiceDiscovery() {
        addressMap = new ConcurrentHashMap<>();
        String loadBalancerName = LoadBalanceSelector.getLoadBalancer();
        loadBalance = ExtensionLoader.of(LoadBalance.class).getExtension(loadBalancerName);
    }

    @Override
    public InetSocketAddress lookupService(RpcRequest rpcRequest) {
        String serviceName = rpcRequest.getFullServiceName();
        List<URL> addressList = getAddressList(serviceName);
        if (addressList.isEmpty()) {
            return null;
        }
        LockUtil.lock(serviceName);
        try {
            return loadBalance.select(addressList, rpcRequest).toInetSocketAddress();
        } finally {
            LockUtil.unlock(serviceName);
        }
    }

    public void addAddress(String serviceName, URL address) {
        List<URL> addressList = getAddressList(serviceName);
        LockUtil.lock(serviceName);
        try {
            addressList.add(address);
        } finally {
            LockUtil.unlock(serviceName);
        }
    }

    protected List<URL> getAddressList(String serviceName) {
        return addressMap.computeIfAbsent(serviceName, k -> new ArrayList<>());
    }

}
